import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import layers, activations, losses, optimizers, metrics
import numpy as np
import os
import sys

np.random.seed(777)
tf.random.set_seed(777)

# 3.	使用Keras实现手写识别
# (1)	数据处理
# ①	加载数据信息（2分）
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
print('x_train', x_train.shape)
print('y_train', y_train.shape)
print('x_test', x_test.shape)
print('y_test', y_test.shape)

# ②	对数据进行必要的预处理操作（3分）
SELECT_RATE = 0.05  # 由于数据集巨大，随机抽取一定比例的样本，正式代码请改成1.0
M, N1, N2 = x_train.shape
M_TRAIN = int(np.ceil(M * SELECT_RATE))
rand_idx = np.random.permutation(M)[:M_TRAIN]
x_train = x_train[rand_idx].reshape(-1, N1*N2)  # ATTENTION Flatten data at first.
y_train = y_train[rand_idx].reshape(-1, 1)
M, N1, N2 = x_test.shape
M_TEST = int(np.ceil(M * SELECT_RATE))
rand_idx = np.random.permutation(M)[:M_TEST]
x_test = x_test[rand_idx].reshape(-1, N1*N2)
y_test = y_test[rand_idx].reshape(-1, 1)
# 特征缩放
x_train = x_train.astype(np.float32) / 255.
x_test = x_test.astype(np.float32) / 255.
# check again
print('x_train', x_train.shape)
print('y_train', y_train.shape)
print('x_test', x_test.shape)
print('y_test', y_test.shape)

# ③	将标签值进行独热编码处理（5分 ）
y_train_oh = keras.utils.to_categorical(y_train)
y_test_oh = keras.utils.to_categorical(y_test)

# (2)	模型处理（每小问5分）
# ①	实现全连接处理，每层神经元数量为256,256,256，256,10
# ②	每层添加dropout处理，失活比例0.3
L1 = 256
L2 = 256
L3 = 256
L4 = 256
N_CLS = 10
dropout = 0.3
model = keras.Sequential([
    layers.Dense(L1, activation=activations.relu, input_shape=(N1 * N2,)),
    layers.Dropout(dropout),
    layers.Dense(L2, activation=activations.relu),
    layers.Dropout(dropout),
    layers.Dense(L3, activation=activations.relu),
    layers.Dropout(dropout),
    layers.Dense(L4, activation=activations.relu),
    layers.Dropout(dropout),
    layers.Dense(N_CLS),
])
model.summary()

# ③	编译并训练数据
ALPHA = 0.001
BATCH_SIZE = 64
N_EPOCHS = 6  # 为了快速演示，这里设置的偏小，正式代码请适量增大
model.compile(
    loss=losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=optimizers.Adam(learning_rate=ALPHA),
    metrics=[metrics.sparse_categorical_accuracy]
)
model.fit(
    x_train,
    y_train,
    batch_size=BATCH_SIZE,
    epochs=N_EPOCHS,
    validation_data=(x_test, y_test),
    validation_batch_size=BATCH_SIZE
)

# ④	打印模型的准确率
model.evaluate(
    x_test,
    y_test,
    batch_size=BATCH_SIZE
)
